{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using TAO EfficientDet (TF2)\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Sample prediction of EfficientDet\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/common/sample.jpg?raw=true\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained model and train an EfficientDet-D0 model on COCO dataset\n",
    "* Evaluate the trained model\n",
    "* Run pruning and finetuning with the trained model\n",
    "* Run inference with the trained model and visualize the result\n",
    "* Export the trained model to a .onnx file for deployment to DeepStream\n",
    "* Run inference on the exported .onnx model to verify deployment using TensorRT\n",
    "\n",
    "At the end of this notebook, you will have generated a trained and optimized `EfficientDet` model\n",
    "which you may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps)\n",
    "or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "This notebook shows an example use case for instance segmentation using the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO Launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2)\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained model](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained model](#head-8)\n",
    "9. [Visualize inferences](#head-9)\n",
    "10. [Deploy](#head-10)\n",
    "11. [Verify the deployed model](#head-11)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/efficientdet_tf2`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/efficientdet_tf2\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/efficientdet_tf2\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/efficientdet_tf2\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "%env LOCAL_PROJECT_DIR=/workspace/tao-experiments/\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"efficientdet_tf2\"\n",
    ")\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/efficientdet_tf2/specs\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit-archive/tlt-30/text/tlt_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "            # Mapping the data directory\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "                \"destination\": \"/workspace/tao-experiments\"\n",
    "            },\n",
    "            # Mapping the specs directory.\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "                \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "            },\n",
    "        ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid()),\n",
    "        \"network\": \"host\"\n",
    "    }\n",
    "}\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After setting up your virtual environment with the above requirements, install TAO pip package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-pyindex\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the COCO dataset for the tutorial. The following script will download COCO dataset automatically and convert it to TFRecords. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR\n",
    "# Download and preprocess data\n",
    "!tao model efficientdet_tf2 run bash $SPECS_DIR/download_coco.sh $DATA_DOWNLOAD_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# convert training data to TFRecords\n",
    "!sed -i \"s|DATADIR|$DATA_DOWNLOAD_DIR|g\" $LOCAL_SPECS_DIR/spec_train.yaml\n",
    "!tao model efficientdet_tf2 dataset_convert -e $SPECS_DIR/spec_train.yaml \\\n",
    "    dataset_convert.results_dir=$DATA_DOWNLOAD_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert validation data to TFRecords\n",
    "!tao model efficientdet_tf2 dataset_convert -e $SPECS_DIR/spec_train.yaml \\\n",
    "    dataset_convert.image_dir=\"$DATA_DOWNLOAD_DIR/raw-data/val2017/\"\\\n",
    "    dataset_convert.annotations_file=\"$DATA_DOWNLOAD_DIR/raw-data/annotations/instances_val2017.json\"\\\n",
    "    dataset_convert.tag='val' \\\n",
    "    dataset_convert.num_shards=32"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the dataset conversion scripts provided in `specs` are intended for the standard COCO dataset. If your data doesn't have `caption` groundtruth or test set, you can modify `download_and_preprocess_coco.sh` and `create_coco_tf_record.py` by commenting out corresponding variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "!ls -l $LOCAL_DATA_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download pretrained model from NGC"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will use NGC CLI to get the pre-trained models. For more details, go to ngc.nvidia.com and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm -f $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/pretrained_efficientdet_tf2:efficientnet_b0*\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_efficientdet_tf2:efficientnet_b0 --dest $LOCAL_EXPERIMENT_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/pretrained_efficientdet_tf2_vefficientnet_b0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Tfrecords for the train datasets\n",
    "    * In order to use the newly generated tfrecords, update the dataset_config parameter in the spec file at `$SPECS_DIR/efficientdet_d0_train.txt` \n",
    "* Note that the learning rate in the spec file is set for 1 GPU training. If you have N gpus, you should multiply LR by N.\n",
    "* \"num_examples_per_epoch\" should be set to the total number of images in the dataset.\n",
    "* Pre-trained models\n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!sed -i \"s|DATADIR|$DATA_DOWNLOAD_DIR|g\" $LOCAL_SPECS_DIR/spec_train.yaml\n",
    "!sed -i \"s|DATADIR|$DATA_DOWNLOAD_DIR|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml\n",
    "!cat $LOCAL_SPECS_DIR/spec_train.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Train an Efficientdet model <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* Evaluation uses COCO metrics. For more info, please refer to: https://cocodataset.org/#detection-eval\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned\n",
    "!sed -i \"s|RESULTSDIR|$USER_EXPERIMENT_DIR/experiment_dir_unpruned|g\" $LOCAL_SPECS_DIR/spec_train.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"For multi-GPU, change num_gpus based on your machine.\")\n",
    "!tao model efficientdet_tf2 train -e $SPECS_DIR/spec_train.yaml num_gpus=$NUM_GPUS \\\n",
    "    results_dir=\"$USER_EXPERIMENT_DIR/experiment_dir_unpruned\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"To resume training from a checkpoint, simply run the same training script. It will pick up from where it's left.\")\n",
    "# !tao model efficientdet_tf2 train -e $SPECS_DIR/spec_train.yaml num_gpus=$NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/train"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Tips:** TAO commands use [Hydra](hydra.cc) to parse the spec file, so we can use Hydra override [syntax](https://hydra.cc/docs/advanced/override_grammar/basic/) to easily change parameters without modifying the spec file.\n",
    "For example, if we want to train our model longer, we can run\n",
    "```\n",
    "!tao model efficientdet_tf2 train -e $SPECS_DIR/spec_train.yaml num_gpus=$NUM_GPUS train.num_epochs=100\n",
    "```\n",
    "To check all the existing parameters, we can add `--info` to the command,\n",
    "```\n",
    "!tao model efficientdet_tf2 train -e $SPECS_DIR/spec_train.yaml --info\n",
    "```\n",
    "\n",
    "The syntax applies to all TAO commands, including dataset_convert, train, evaluate, prune, inference and export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last checkpoints\n",
    "last_checkpoint = ''\n",
    "for f in os.listdir(os.path.join(os.environ[\"LOCAL_EXPERIMENT_DIR\"],'experiment_dir_unpruned', 'train')):\n",
    "    if f.startswith('efficientdet-d'):\n",
    "        last_checkpoint = last_checkpoint if last_checkpoint > f else f\n",
    "print(f'Last checkpoint: {last_checkpoint}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set LAST_CHECKPOINT in the spec file\n",
    "%env LAST_CHECKPOINT={last_checkpoint}\n",
    "!sed -i \"s|EVALMODEL|$USER_EXPERIMENT_DIR/experiment_dir_unpruned/train/$LAST_CHECKPOINT|g\" $LOCAL_SPECS_DIR/spec_train.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model efficientdet_tf2 evaluate -e $SPECS_DIR/spec_train.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune <a class=\"anchor\" id=\"head-6\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Specify pre-trained model\n",
    "- Equalization criterion\n",
    "- Threshold for pruning.\n",
    "- Output directory to store the model\n",
    "\n",
    "Usually, you just need to adjust -pth (threshold) for accuracy and model size trade off. Higher pth gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold value depends on the dataset and the model. 0.5 in the block below is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Notes\n",
    "print(\"The pruned model by default will be saved under `$LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/prune/`\")\n",
    "print(\"To change the saving directory, please update `prune.results_dir`\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model efficientdet_tf2 prune -e $SPECS_DIR/spec_train.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/prune"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note** that you should retrain the pruned model first, as it cannot be directly used for evaluation or inference. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Model needs to be re-trained to bring back accuracy after pruning\n",
    "- Specify re-training specification\n",
    "- WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain\n",
    "!sed -i \"s|RESULTSDIR|$USER_EXPERIMENT_DIR/experiment_dir_retrain|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml\n",
    "!sed -i \"s|PRUNEDMODEL|$USER_EXPERIMENT_DIR/experiment_dir_unpruned/prune/model_th=0.5_eq=union.tlt|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model efficientdet_tf2 train -e $SPECS_DIR/spec_retrain.yaml num_gpus=$NUM_GPUS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Quantization aware training (QAT) with the pruned model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat\n",
    "!sed -i \"s|DATADIR|$DATA_DOWNLOAD_DIR|g\" $LOCAL_SPECS_DIR/spec_retrain_qat.yaml\n",
    "!sed -i \"s|RESULTSDIR|$USER_EXPERIMENT_DIR/experiment_dir_retrain_qat|g\" $LOCAL_SPECS_DIR/spec_retrain_qat.yaml\n",
    "!sed -i \"s|PRUNEDMODEL|$USER_EXPERIMENT_DIR/experiment_dir_unpruned/prune/model_th=0.5_eq=union.tlt|g\" $LOCAL_SPECS_DIR/spec_retrain_qat.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model efficientdet_tf2 train -e $SPECS_DIR/spec_retrain_qat.yaml num_gpus=$NUM_GPUS"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last step of saved checkpoints\n",
    "last_checkpoint = ''\n",
    "for f in os.listdir(os.path.join(os.environ[\"LOCAL_EXPERIMENT_DIR\"],'experiment_dir_retrain', 'train')):\n",
    "    if f.startswith('efficientdet-d'):\n",
    "        last_checkpoint = last_checkpoint if last_checkpoint > f else f\n",
    "print(f'Last checkpoint: {last_checkpoint}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set LAST_CHECKPOINT in the spec file\n",
    "%env LAST_CHECKPOINT={last_checkpoint}\n",
    "!sed -i \"s|EVALMODEL|$USER_EXPERIMENT_DIR/experiment_dir_retrain/train/$LAST_CHECKPOINT|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model efficientdet_tf2 evaluate -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `infer` tool to generate inferences on the trained models and visualize the results. The `infer` tool produces annotated image outputs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy some test images\n",
    "!mkdir -p $LOCAL_DATA_DIR/test_samples\n",
    "!cp $LOCAL_DATA_DIR/raw-data/test2017/0000000000* $LOCAL_DATA_DIR/test_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tao model efficientdet_tf2 inference -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# install deps\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Visualizing the sample images.\n",
    "OUTPUT_PATH = 'experiment_dir_retrain/inference' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 2 # number of columns in the visualizer grid.\n",
    "IMAGES = 4 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Deploy! <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!rm -rf $LOCAL_EXPERIMENT_DIR/export\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export\n",
    "\n",
    "# Generate .onnx file using tao container\n",
    "!sed -i \"s|EXPORTDIR|$USER_EXPERIMENT_DIR/export|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml\n",
    "!tao model efficientdet_tf2 export -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if onnx model is correctly saved.\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/export"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference.\n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's.\n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson and refer to [here](https://docs.nvidia.com/tao/tao-toolkit-archive/tao-30-2108/text/tensorrt.html#installing-the-tao-converter) for more details.\n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP32).\n",
    "!tao deploy efficientdet_tf2 gen_trt_engine -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (INT8).\n",
    "!sed -i \"s|fp32|int8|g\" $LOCAL_SPECS_DIR/spec_retrain.yaml\n",
    "!tao deploy efficientdet_tf2 gen_trt_engine -e $SPECS_DIR/spec_retrain.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported models:')\n",
    "print('------------')\n",
    "!ls -lth $LOCAL_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get the last QAT checkpoints\n",
    "last_checkpoint = ''\n",
    "for f in os.listdir(os.path.join(os.environ[\"LOCAL_EXPERIMENT_DIR\"],'experiment_dir_retrain_qat', 'train')):\n",
    "    if f.startswith('efficientdet-d'):\n",
    "        last_checkpoint = last_checkpoint if last_checkpoint > f else f\n",
    "print(f'Last checkpoint: {last_checkpoint}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set LAST_CHECKPOINT in the spec file\n",
    "%env LAST_CHECKPOINT={last_checkpoint}\n",
    "!sed -i \"s|EVALMODEL|$USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/train/$LAST_CHECKPOINT|g\" $LOCAL_SPECS_DIR/spec_retrain_qat.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/export_qat\n",
    "# Export in QAT mode. \n",
    "!sed -i \"s|EXPORTDIR|$USER_EXPERIMENT_DIR/export_qat|g\" $LOCAL_SPECS_DIR/spec_retrain_qat.yaml\n",
    "!tao model efficientdet_tf2 export -e $SPECS_DIR/spec_retrain_qat.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert QAT to TRT engine\n",
    "!tao deploy efficientdet_tf2 gen_trt_engine -e $SPECS_DIR/spec_retrain_qat.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check if onnx model is correctly saved.\n",
    "!ls -l $LOCAL_EXPERIMENT_DIR/export_qat"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify the deployed model <a class=\"anchor\" id=\"head-11\"></a>\n",
    "\n",
    "Verify the converted engine by visualizing TensorRT inferences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on a dir of images\n",
    "!tao deploy efficientdet_tf2 inference -e $SPECS_DIR/spec_retrain.yaml \\\n",
    "    inference.results_dir=$USER_EXPERIMENT_DIR/export"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -l $LOCAL_EXPERIMENT_DIR/export/"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  },
  "vscode": {
   "interpreter": {
    "hash": "767d51c1340bd893661ea55ea3124f6de3c7a262a8b4abca0554b478b1e2ff90"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
